#! env python

from posixpath import split
import re
import subprocess as subp
from typing import AnyStr, List

def parse_ptxas_output(lines: List[AnyStr]):
    target_code = lines[0].split()[-1].strip("'")
    
    fn_decorated = lines[1].split()[-1]

    stack_info = lines[2].split(",")
    stack_frames = int(stack_info[0].split()[0])
    spilt_loads = int(stack_info[2].split()[0])
    split_stores = int(stack_info[1].split()[0])

    reg_info = lines[3].split(',')
    regs = int(reg_info[0].split()[1])
    cmem = int(reg_info[-1].split()[0])
    smem = int(reg_info[1].split()[0]) if len(reg_info) > 2 else 0

    fn_name = subp.check_output(['c++filt', fn_decorated])
    fn_name = str(fn_name, encoding='utf-8').strip()
    
    return f"Arch {target_code}-Kernel:{fn_name}:\n"\
           f"res=[{regs} reg, {smem} smem, {cmem} cmem],"\
           f"split=[{stack_frames} local, {spilt_loads}/{split_stores} splitL/S]"


if __name__ == '__main__':
    import sys
    with open(sys.argv[1], 'r') as f:
        line = f.readline()
        while re.match('ptxas info', line) is None:
            line = f.readline()
        
        while line:
            line = f.readline()
            if not re.search("Compiling entry function", line) is None:
                line = line.replace('ptxas info    : ','').strip()
                args = [line]
                for i in range(3):
                    args.append(f.readline().replace('ptxas info    : ','').strip())
                print(parse_ptxas_output(args))
            
            
        